#include "MLRegress_Linear.h"
#include <stdio.h>
#include <gsl/gsl_multifit.h>
_MYLABLIB_BEGIN

CMLRegress_Linear::CMLRegress_Linear(void)
{
}

CMLRegress_Linear::~CMLRegress_Linear(void)
{
}

/** 
	@Name: MultiVariant linear regress function
	@Desc:
	Suppose the linear regression model is : Y = AX
	Y is the data points of the last element of the input parameter 'dataSets'
	X is the data points of other elements( except the last) of 'dataSets'.
	A is Output parameter 'coefficients'. The vector of coefficients is make of every dimension's coefficients
*/
int CMLRegress_Linear::MultiVariantRegress(int nDimension, const std::vector<CMLDataSet>& dataSets, std::vector<CMLDataPoint>& coefficients)
{
	int i,j;
	gsl_matrix* X;
	gsl_matrix* COV;
	gsl_vector* Y;
	gsl_vector* C;
	gsl_multifit_linear_workspace* work;
	double chisq = 0.0;
	int p = (int)dataSets.size() - 1;

	// Initialize the matrix and vectors
	int nCount = dataSets[0].Count();
	work = gsl_multifit_linear_alloc((size_t)nCount, p);
	X = gsl_matrix_alloc((size_t)nCount, (size_t)p);
	COV = gsl_matrix_alloc((size_t)p, (size_t)p);
	Y = gsl_vector_alloc((size_t)nCount);
	C = gsl_vector_alloc((size_t)p);

	// Clear the coefficients
	coefficients.clear();
	coefficients.resize((size_t)p);

	// Calculate the coefficients for each dimension
	for (int d=0; d<nDimension; d++)
	{
		// Set the X matrix
		for (i=0; i<nCount; i++)
		{
			for (j=0; j<p; j++)
				gsl_matrix_set(X, i, j, dataSets[j].Get(i)[d]);
		}
		// Set the Y vector
		for (i=0; i<nCount; i++)
			gsl_vector_set(Y, i, dataSets[p].Get(i)[d]);

		// Use Least-squares to linear regress
		gsl_multifit_linear(X, Y, C, COV, &chisq, work);

		// Extract the C to cofficients
		CMLDataPoint& coe = coefficients[d];
		coe.resize(p);
		for (i=0; i<p; i++)
			coe[i] = gsl_vector_get(C, (size_t)i);
	}

	// Clean the matrix and vectors
	gsl_matrix_free(X);
	gsl_matrix_free(COV);
	gsl_vector_free(Y);
	gsl_vector_free(C);
	gsl_multifit_linear_free(work);
	return 0;
}

/** 
	@Name: AutoRegress function
	@Desc:
	Suppose the linear regression model is : X(t) = A(1)X(t-1) + A(2)X(t-2) + ... + A(p-1)X(t-(p-1)) + A(p)
	X(1), X(2), ... , X(n) is the input parameter 'dataSet'
	p is the size of autoregress parameters;
	(A(1),...,A(p)) is Output parameter 'coefficients'. The vector of coefficients is make of every dimension's coefficients
*/
int CMLRegress_Linear::AR_Regress(int nDimension, const CMLDataSet& dataSet, int p, std::vector<CMLDataPoint>& coefficients,  CMLDataPoint& chisqs)
{
	int i,j;
	gsl_matrix* X;
	gsl_matrix* COV;
	gsl_vector* Y;
	gsl_vector* C;
	gsl_multifit_linear_workspace* work;
	double chisq = 0.0;

	ML_ASSERT(dataSet.Count() >= p);

	// Initialize the matrix and vectors
	int nCount = dataSet.Count();
	int nUseCount = nCount - (p-1);
	work = gsl_multifit_linear_alloc((size_t)nUseCount, p);
	X = gsl_matrix_alloc((size_t)nUseCount, (size_t)p);
	COV = gsl_matrix_alloc((size_t)p, (size_t)p);
	Y = gsl_vector_alloc((size_t)nUseCount);
	C = gsl_vector_alloc((size_t)p);

	// Clear the coefficients
	coefficients.clear();
	coefficients.resize((size_t)nDimension);
	chisqs.resize(nDimension);

	// Calculate the coefficients for each dimension
	for (int d=0; d<nDimension; d++)
	{
		// Set the X matrix
		for (i=p-1; i<nCount; i++)
		{
			gsl_matrix_set(X, i-p+1, 0, 1.0);
			for (j=1; j<p; j++)
				gsl_matrix_set(X, i-p+1, j, dataSet.Get(i-j)[d]);
		}
		// Set the Y vector
		for (i=p-1; i<nCount; i++)
			gsl_vector_set(Y, i-p+1, dataSet.Get(i)[d]);

		// Use Least-squares to linear regress
		gsl_multifit_linear(X, Y, C, COV, &chisq, work);

		// Extract the C to cofficients
		CMLDataPoint& coe = coefficients[d];
		coe.resize(p);
		for (i=0; i<p; i++)
			coe[i] = gsl_vector_get(C, (size_t)i);

		chisqs[d] = chisq;
	}

	// Clean the matrix and vectors
	gsl_matrix_free(X);
	gsl_matrix_free(COV);
	gsl_vector_free(Y);
	gsl_vector_free(C);
	gsl_multifit_linear_free(work);
	return 0;
}

int CMLRegress_Linear::AR_Predict(int nDimension, CMLDataSet& predicts, int p, std::vector<CMLDataPoint>& coefficients, int nPredictLen)
{
	ML_ASSERT(predicts.Count() >= p-1);

	CMLDataPoint predictPoint(nDimension);
	for (int i=p-1; i<nPredictLen+p-1; i++)
	{
		for (int d=0; d<nDimension; d++)
		{
			predictPoint[d] = coefficients[d][0];
			for (int j=1; j<p;j++)
				predictPoint[d] += coefficients[d][j]*predicts.Get(i-p+j)[d];
		}
		predicts.Insert(predictPoint);
	}
	return 0;
}

_MYLABLIB_END




